# _*_ coding: utf-8 _*_
# ---------------------------
# @创建作者：ZQQ
# @创建时间：2025/8/28 14:50
# @说明：
# ---------------------------
import os
import time
import zipfile
from pathlib import Path

import requests


class ModelDownloader:
    def __init__(self, url, target_folder, chunk_size=8192, max_retries=3):
        """
        初始化下载器实例
        :param url: 文件下载URL
        :param target_folder: 目标解压文件夹路径（如 'server/100000001-8775/weights'）
        :param chunk_size: 分块大小(字节)
        :param max_retries: 最大重试次数
        """
        # 目标解压文件夹路径
        self.extract_path = Path.cwd() / target_folder
        print(f'目标解压目录:{self.extract_path}')
        if not os.path.exists(self.extract_path):
            # 自动创建多级目录
            self.extract_path.mkdir(parents=True, exist_ok=True)

        # 从URL中提取文件名（处理URL参数）
        self.url = url
        url_filename = url.split('/')[-1].split('?')[0]  # 处理URL参数
        
        # 判断文件类型（不区分大小写）
        url_filename_lower = url_filename.lower()
        is_zip_file = url_filename_lower.endswith('.zip')
        is_joblib_file = url_filename_lower.endswith('.joblib')
        
        # 逻辑：如果是.zip文件，使用zip逻辑；否则（包括.joblib和其他情况）都当作joblib处理
        self.is_joblib_file = not is_zip_file  # 不是zip文件就当作joblib处理
        
        print(f'🔍 检测文件类型: URL文件名={url_filename}, 是否为zip文件={is_zip_file}, 是否为joblib文件={is_joblib_file}, 最终处理方式={"joblib" if self.is_joblib_file else "zip"}')
        
        if self.is_joblib_file:
            self.model_filename = 'model.joblib'
            self.model_path = self.extract_path / self.model_filename
            self.temp_path = self.model_path.with_suffix('.tmp')  # 临时文件路径
            self.zip_path = None  # 不需要zip路径
            print(f'✅ 识别为joblib文件（或默认joblib），将直接下载到: {self.model_path}')
        else:
            # 如果是 zip 文件，使用原有逻辑
            zip_filename = url_filename if is_zip_file else "model.zip"
            # 生成保存路径(下载压缩包，保存在目标文件夹的父目录)
            self.deploy_dir = self.extract_path.parent
            self.zip_path = self.deploy_dir / zip_filename
            self.temp_path = self.zip_path.with_suffix('.tmp')  # 临时文件路径
            self.model_path = None  # 不需要model路径
            print(f'✅ 识别为zip文件，将下载压缩包到: {self.zip_path}')
        
        self.chunk_size = chunk_size
        self.max_retries = max_retries

    def _extract_zip(self) -> str:
        """
        解压zip文件，忽略压缩包内的文件夹名，直接解压到目标文件夹
        :return: 解压后的目标文件夹路径
        """
        try:
            # 确保解压目录存在
            if not self.extract_path.exists():
                self.extract_path.mkdir(parents=True, exist_ok=True)
            
            with zipfile.ZipFile(self.zip_path, 'r') as zip_ref:
                # 获取zip文件中所有文件列表
                namelist = zip_ref.namelist()
                
                # 找到所有非空文件路径
                non_empty_names = [name for name in namelist if name.strip() and not name.endswith('/')]
                
                if not non_empty_names:
                    # 如果压缩包为空，直接返回
                    print('⚠️ 压缩包为空')
                    return ""
                
                # 找到所有文件的第一层目录（如果有的话）
                first_level_dirs = set()
                for name in non_empty_names:
                    parts = name.split('/')
                    if len(parts) > 1 and parts[0]:
                        first_level_dirs.add(parts[0])
                
                # 如果所有文件都在同一个第一层目录下，跳过这一层（忽略压缩包内的文件夹名）
                if len(first_level_dirs) == 1:
                    # 所有文件都在同一个第一层目录下，跳过这一层
                    skip_prefix = list(first_level_dirs)[0] + '/'
                    for member in zip_ref.infolist():
                        if member.filename.startswith(skip_prefix):
                            # 移除第一层目录前缀
                            original_filename = member.filename
                            member.filename = member.filename[len(skip_prefix):]
                            if member.filename:  # 跳过空的目录
                                # 确保目标路径安全（防止路径遍历攻击）
                                target_path = self.extract_path / member.filename
                                try:
                                    # 使用 resolve() 获取绝对路径，然后检查是否以 extract_path 开头
                                    extract_abs = self.extract_path.resolve()
                                    target_abs = target_path.resolve()
                                    if str(target_abs).startswith(str(extract_abs) + os.sep) or str(target_abs) == str(extract_abs):
                                        zip_ref.extract(member, self.extract_path)
                                except (ValueError, OSError):
                                    # 路径不安全或出错，跳过
                                    pass
                            member.filename = original_filename  # 恢复原始文件名，避免影响后续处理
                else:
                    # 文件不在同一个第一层目录下，或者没有第一层目录，直接解压所有内容
                    for member in zip_ref.infolist():
                        if member.filename and not member.filename.endswith('/'):
                            # 确保目标路径安全（防止路径遍历攻击）
                            target_path = self.extract_path / member.filename
                            try:
                                # 使用 resolve() 获取绝对路径，然后检查是否以 extract_path 开头
                                extract_abs = self.extract_path.resolve()
                                target_abs = target_path.resolve()
                                if str(target_abs).startswith(str(extract_abs) + os.sep) or str(target_abs) == str(extract_abs):
                                    zip_ref.extract(member, self.extract_path)
                            except (ValueError, OSError):
                                # 路径不安全或出错，跳过
                                pass
                print('✅ 模型压缩包解压完成！')
            
            # 删除压缩包以节省空间
            self.zip_path.unlink()
            print('🗑️ 压缩包已删除')
            
            # 清理临时文件（如果存在）
            if self.temp_path.exists():
                try:
                    self.temp_path.unlink()
                except (OSError, ValueError):
                    pass
            
            # 返回目标文件夹路径
            if self.extract_path.exists() and self.extract_path.is_dir():
                return str(self.extract_path.resolve())
            
            print(f'❌ 未找到目标文件夹: {self.extract_path}')
            return ""
        except Exception as e:
            print(f'❌ 解压失败: {e}')
            return ""

    def _download_joblib(self) -> str:
        """
        直接下载 .joblib 文件到目标文件夹
        :return: 目标文件夹绝对路径（失败返回空字符串）
        """
        try:
            # 检查文件是否已存在
            if self.model_path.exists():
                print(f'✅ 模型文件已存在: {self.model_path}')
                return str(self.extract_path.resolve())
            
            # 断点续传准备
            existing_size = self.temp_path.stat().st_size if self.temp_path.exists() else 0

            # 重试下载逻辑（指数退避）
            for attempt in range(self.max_retries):
                try:
                    headers = {'Range': f'bytes={existing_size}-'} if existing_size else {}
                    
                    # 设置连接超时和读取超时（连接5秒，读取30秒）
                    timeout = (5, 30)  # (connect_timeout, read_timeout)

                    with requests.get(self.url, headers=headers, stream=True, timeout=timeout) as r:
                        r.raise_for_status()

                        # 流式写入文件，添加超时保护
                        mode = 'ab' if existing_size else 'wb'
                        download_start_time = time.time()
                        last_data_time = time.time()
                        chunk_timeout = 60  # 每个chunk之间最多等待60秒
                        total_timeout = 300  # 总下载时间最多5分钟
                        
                        with open(self.temp_path, mode) as f:
                            for chunk in r.iter_content(chunk_size=self.chunk_size):
                                current_time = time.time()
                                
                                # 检查总超时
                                if current_time - download_start_time > total_timeout:
                                    raise requests.exceptions.Timeout(f'下载总超时：超过{total_timeout}秒')
                                
                                if chunk:  # 过滤keep-alive空块
                                    f.write(chunk)
                                    last_data_time = current_time
                                else:
                                    # 检查数据接收超时（即使收到空块，也应该定期有响应）
                                    if current_time - last_data_time > chunk_timeout:
                                        raise requests.exceptions.Timeout(f'下载超时：超过{chunk_timeout}秒未收到数据')
                            print('📦 模型文件下载完成！')

                    # 原子操作：临时文件转正式文件
                    self.temp_path.rename(self.model_path)
                    
                    # 返回目标文件夹路径
                    if self.extract_path.exists() and self.extract_path.is_dir():
                        return str(self.extract_path.resolve())
                    
                    print(f'❌ 未找到目标文件夹: {self.extract_path}')
                    return ""
                    
                except (requests.RequestException, IOError) as e:
                    print(f'🔄 尝试 {attempt + 1}/{self.max_retries} 失败: {e}')
                    if attempt < self.max_retries - 1:
                        time.sleep(2 ** attempt)  # 指数退避等待
            return ""  # 所有重试失败
        except Exception as e:
            print(f'❌ 下载joblib文件失败: {e}')
            return ""

    @property
    def download(self) -> str:
        """
        执行下载逻辑（自动跳过已下载文件）
        支持两种模式：
        1. .joblib 文件：直接下载到目标文件夹
        2. .zip 压缩包：下载后解压到目标文件夹
        :return: 目标文件夹绝对路径（失败返回空字符串）
        """
        # 检查目标文件夹是否存在
        if not self.extract_path.exists():
            self.extract_path.mkdir(parents=True, exist_ok=True)
        
        # 如果是 .joblib 文件，使用直接下载方式
        if self.is_joblib_file:
            print(f'📥 使用joblib文件下载模式')
            # 检查文件是否已存在
            if self.model_path and self.model_path.exists():
                print(f'✅ 模型文件已存在: {self.model_path}')
                return str(self.extract_path.resolve())
            result = self._download_joblib()
            if result:
                print(f'✅ joblib文件下载成功: {result}')
            else:
                print(f'❌ joblib文件下载失败')
            return result
        
        # 如果是 zip 文件，使用原有逻辑
        print(f'📥 使用zip文件下载模式')
        # 检查目标文件夹是否存在且有内容
        if self.extract_path.exists() and self.extract_path.is_dir():
            # 检查文件夹是否有内容（至少有一个文件或子目录）
            if any(self.extract_path.iterdir()):
                print(f'✅ 目标文件夹已存在: {self.extract_path}')
                return str(self.extract_path.resolve())

        # 检查压缩包是否已存在（确保 zip_path 不为 None）
        if self.zip_path and self.zip_path.exists():
            print(f'✅ 压缩包已存在，直接解压: {self.zip_path}')
            return self._extract_zip()

        # 断点续传准备
        existing_size = self.temp_path.stat().st_size if self.temp_path.exists() else 0

        # 重试下载逻辑（指数退避）
        for attempt in range(self.max_retries):
            try:
                headers = {'Range': f'bytes={existing_size}-'} if existing_size else {}
                
                # 设置连接超时和读取超时（连接5秒，读取30秒）
                timeout = (5, 30)  # (connect_timeout, read_timeout)

                with requests.get(self.url, headers=headers, stream=True, timeout=timeout) as r:
                    r.raise_for_status()

                    # 流式写入文件，添加超时保护
                    mode = 'ab' if existing_size else 'wb'
                    download_start_time = time.time()
                    last_data_time = time.time()
                    chunk_timeout = 60  # 每个chunk之间最多等待60秒
                    total_timeout = 300  # 总下载时间最多5分钟
                    
                    with open(self.temp_path, mode) as f:
                        for chunk in r.iter_content(chunk_size=self.chunk_size):
                            current_time = time.time()
                            
                            # 检查总超时
                            if current_time - download_start_time > total_timeout:
                                raise requests.exceptions.Timeout(f'下载总超时：超过{total_timeout}秒')
                            
                            if chunk:  # 过滤keep-alive空块
                                f.write(chunk)
                                last_data_time = current_time
                            else:
                                # 检查数据接收超时（即使收到空块，也应该定期有响应）
                                if current_time - last_data_time > chunk_timeout:
                                    raise requests.exceptions.Timeout(f'下载超时：超过{chunk_timeout}秒未收到数据')
                        print('📦 模型压缩包下载完成！')

                # 原子操作：临时文件转正式文件
                self.temp_path.rename(self.zip_path)
                
                # 解压文件
                return self._extract_zip()

            except (requests.RequestException, IOError) as e:
                print(f'🔄 尝试 {attempt + 1}/{self.max_retries} 失败: {e}')
                if attempt < self.max_retries - 1:
                    time.sleep(2 ** attempt)  # 指数退避等待
        return ""  # 所有重试失败


if __name__ == "__main__":
    # 使用说明：
    # 1. 压缩上一级目录的 weights 文件夹为 zip 文件（例如：cd /data1/code && zip -r weights.zip weights/）
    # 2. 启动一个简单的 HTTP 服务器（例如：cd /data1/code && python -m http.server 8888）
    # 3. 将生成的 URL 填入下面的 download_url 变量
    # 例如：download_url = "http://localhost:8888/weights.zip"
    
    # 下载 URL（请替换为实际的 URL）
    download_url = "http://localhost:8888/we.zip"
    
    # 使用 ModelDownloader 下载并解压到指定文件夹
    # 传入目标文件夹路径，会解压到该文件夹内（忽略压缩包内的文件夹名）
    res = ModelDownloader(download_url, 'temp').download
    print(f'✅ weights文件夹路径:\n{res}')
